#!/usr/bin/env python

from websocket import WebSocketApp
from time import sleep
from json import loads, dumps
from json.decoder import JSONDecodeError
from multiprocessing import Pool
from common import make_verify_sign, loop, get_center_md5, pidinfo, redis, hwinfo
from config import WS_URL, BS_RDS_PERFIX, KB_RDS_KEY, KB_LOCK_VAL
from net import center_post
from log import out, err


class Center:

    def __init__(self, max_pool=16, **kwargs):
        self.pool = Pool(max_pool)
        self.ID = get_center_md5()
        self.key = None
        self.uuid = None
        self.ws = self.wsInit()
        self.rds = self.cacheInit()

    def quit(self):
        self.pool.close()
        self.pool.join()
        out('Center Quit...')

    def cacheInit(self):
        rds = redis()
        #  删除所有浏览器指令相关缓存
        bskeys = rds.keys('%s_*' % BS_RDS_PERFIX)
        if bskeys:
            rds.delete(*bskeys)
        #  设置键盘锁初始值
        if not rds.exists(KB_RDS_KEY):
            rds.delete(KB_RDS_KEY)
        rds.rpush(KB_RDS_KEY, KB_LOCK_VAL)
        return rds

    def wsInit(self):
        #  websocket.enableTrace(True)
        ws = WebSocketApp(WS_URL, keep_running=True)
        ws.on_open = self.register
        ws.on_message = self.wsMessage
        ws.on_error = self.wsError
        ws.on_pong = self.wsPong
        ws.on_close = self.wsClose
        return ws

    def register(self):
        try:
            r = center_post(**hwinfo())
            if r['code']:
                out('准备重新连接发送....')
                sleep(2)
                return self.register()

            out('Center Register: ', r)
            self.key = r['request']['sign']
            self.uuid = r['uuid']
            self.options = r['options']

            self.send('login', uuid=self.uuid, key=self.key)
        except:
            err('基站注册失败: ', stack=True)

    def send(self, msg_type, **kwargs):
        msg = {}
        msg['type'] = msg_type
        msg['data'] = kwargs
        if kwargs.get('out', True):
            out('WebSocket Send: ', msg)
        self.ws.send(dumps(msg))

    def verify(self, msg):
        """
        签名验证, 如果有sign和data字段才签名验证, 否则不认证
        """
        if msg.get('sign', None) and msg.get('data', None):
            return make_verify_sign(loads(msg['data']), msg['sign'],
                                    'WebSocket', self.key)
        return True

    def taskLogin(self, data):
        """
        WebSocket登录消息处理方法
        """
        if data != 'ok':
            self.register()

    def taskErrorCallback(self, value):
        """
        进程池任务调用错误的回调函数, 并非任务执行错误
        """
        err('apply_async调用任务错误(%s): ' % pidinfo(), value)

    def execTask(self, name, data):
        """
        执行任务
            查找task目录加载对应任务模块方法
            使用进程池执行任务模块方法
            如果task目录查找失败, 则在本类查找task开头方法并且执行
        """
        try:
            module = __import__("task.%s" % name)
            out('Task import: ', module)
            module = getattr(module, name)
            out('Task import Type: ', module)
            entry = getattr(module, 'Task')
            out('Task import Entry: ', entry)
            return self.pool.apply_async(entry(self.uuid, self.key, data,
                                               options=self.options),
                                         error_callback=self.taskErrorCallback)
        except ImportError as e:
            entry = getattr(self, 'task%s' % name.title(), None)
            if entry is None:
                return err('未匹配的任务类型: %s(%s)' % (name, data), stack=True)
            return entry(data)

    def wsMessage(self, msg):
        """
        消息处理器
        """
        try:
            msg = loads(msg)
            out('Task Msg: ', msg)
            if not self.verify(msg):
                return err("签名验证失败!")
            self.execTask(msg.get('type'), msg.get('data', None))
            out('Task Entry End')
        except JSONDecodeError:
            err('Task JSON Data Parse Error!')
        except AttributeError as e:
            err("Task Not Found: %s" % e, stack=True)
        except:
            err("WS Error: ", stack=True)

    def wsPong(self, *args, **kwargs):
        self.send("ping", out=False)

    def wsError(self, e):
        """
        错误处理器
        """
        #  错误回溯
        e = e.with_traceback(None)
        err(e)

    def wsClose(self):
        """
        关闭连接处理器
        """
        out('### Websocket closed ###')

    def run(self):
        out('### WebSocket Run Begin ###')
        loop(self.ws.run_forever, delay=2, ping_interval=30, ping_timeout=15)
        out('### WebSocket Run End ###')

